/*
 *   Copyright 2021 zzh
 *
 *   Licensed under the Apache License, Version 2.0 (the "License");
 *   you may not use this file except in compliance with the License.
 *   You may obtain a copy of the License at
 *
 *       http://www.apache.org/licenses/LICENSE-2.0
 *
 *   Unless required by applicable law or agreed to in writing, software
 *   distributed under the License is distributed on an "AS IS" BASIS,
 *   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *   See the License for the specific language governing permissions and
 *   limitations under the License.
 *
 */

package com.aduib.boot.common.buffer;

import com.aduib.boot.common.concurrent.JctoolsHelper;
import com.aduib.boot.common.lang.ObjectPool;
import com.aduib.boot.common.log.LOG;

import java.util.ArrayList;
import java.util.List;
import java.util.Queue;
import java.util.concurrent.atomic.AtomicBoolean;

import static com.aduib.boot.common.buffer.BufferPoolArena.bufferArea.Normal;
import static com.aduib.boot.common.buffer.BufferPoolArena.bufferArea.Small;

/**
 * @description: PoolThreadCache
 * @author: zzh
 * @date: 2021/12/2 11:35
 */
public class PoolThreadCache {

  protected static final int INTEGER_SIZE_MINUS_ONE = Integer.SIZE - 1;

  final BufferPool bufferPool;

  private final MemoryRegionCache[] normalDirectCaches;

  private final MemoryRegionCache[] smallSubPageDirectCaches;

  // Used for bitshifting when calculate the index of normal caches later
  private final int numShiftsNormalDirect;

  private final int freeSweepAllocationThreshold;

  private final AtomicBoolean freed = new AtomicBoolean();

  private int allocations;

  public PoolThreadCache(
      BufferPool bufferPool,
      int cacheSize,
      int maxCachedBufferCapacity,
      int freeSweepAllocationThreshold) {
    this.bufferPool = bufferPool;
    numShiftsNormalDirect = log2(bufferPool.getPageSize());
    this.freeSweepAllocationThreshold = freeSweepAllocationThreshold;
    normalDirectCaches =
        createNormalCaches2(cacheSize, maxCachedBufferCapacity, (BufferPoolArena) bufferPool);
    smallSubPageDirectCaches =
        createSubPageCaches(cacheSize, ((BufferPoolArena) bufferPool).numSmallSubpagePools);
    ((BufferPoolArena) bufferPool).numThreadCaches.getAndIncrement();
  }

  private static MemoryRegionCache[] createSubPageCaches(int cacheSize, int numCaches) {
    if (cacheSize > 0 && numCaches > 0) {
      @SuppressWarnings("unchecked")
      MemoryRegionCache[] cache = new MemoryRegionCache[numCaches];
      for (int i = 0; i < cache.length; i++) {
        // TODO: maybe use cacheSize / cache.length
        cache[i] = new SubPageMemoryRegionCache(cacheSize);
      }
      return cache;
    } else {
      return null;
    }
  }

  private static MemoryRegionCache[] createNormalCaches(
      int cacheSize, int maxCachedBufferCapacity, BufferPool area) {
    if (cacheSize > 0 && maxCachedBufferCapacity > 0) {
      int max = Math.min(area.getChunkSize(), maxCachedBufferCapacity);
      int arraySize = Math.max(1, log2(max / area.getPageSize()) + 1);

      @SuppressWarnings("unchecked")
      MemoryRegionCache[] cache = new MemoryRegionCache[arraySize];
      for (int i = 0; i < cache.length; i++) {
        cache[i] = new NormalMemoryRegionCache(cacheSize);
      }
      return cache;
    } else {
      return null;
    }
  }

  private static MemoryRegionCache[] createNormalCaches2(
      int cacheSize, int maxCachedBufferCapacity, BufferPoolArena area) {
    if (cacheSize > 0 && maxCachedBufferCapacity > 0) {
      int max = Math.min(area.getChunkSize(), maxCachedBufferCapacity);

      // Create as many normal caches as we support based on how many sizeIdx we have and what the
      // upper
      // bound is that we want to cache in general.
      List<MemoryRegionCache> cache = new ArrayList<>();
      for (int idx = area.nSubpages; idx < area.nSizes && area.sizeIdx2size(idx) <= max; idx++) {
        cache.add(new NormalMemoryRegionCache(cacheSize));
      }
      return cache.toArray(new MemoryRegionCache[0]);
    } else {
      return null;
    }
  }

  private static void trim(MemoryRegionCache[] caches) {
    if (caches == null) {
      return;
    }
    for (MemoryRegionCache c : caches) {
      trim(c);
    }
  }

  // val > 0
  private static int log2(int val) {
    return INTEGER_SIZE_MINUS_ONE - Integer.numberOfLeadingZeros(val);
  }

  private static void trim(MemoryRegionCache cache) {
    if (cache == null) {
      return;
    }
    cache.trim();
  }

  private static MemoryRegionCache cache(MemoryRegionCache[] cache, int sizeIdx) {
    if (cache == null || sizeIdx > cache.length - 1) {
      return null;
    }
    return cache[sizeIdx];
  }

  /**
   * Fast method of finding the next power of 2 greater than or equal to the supplied value.
   *
   * <p>If the value is {@code <= 0} then 1 will be returned. This method is not suitable for {@link
   * Integer#MIN_VALUE} or numbers greater than 2^30.
   *
   * @param value from which to search for next power of 2
   * @return The next power of 2 or the value itself if it is a power of 2
   */
  public static int findNextPositivePowerOfTwo(final int value) {
    assert value > Integer.MIN_VALUE && value < 0x40000000;
    return 1 << (32 - Integer.numberOfLeadingZeros(value - 1));
  }

  /**
   * Fast method of finding the next power of 2 greater than or equal to the supplied value.
   *
   * <p>This method will do runtime bounds checking and call {@link
   * #findNextPositivePowerOfTwo(int)} if within a valid range.
   *
   * @param value from which to search for next power of 2
   * @return The next power of 2 or the value itself if it is a power of 2.
   *     <p>Special cases for return values are as follows:
   *     <ul>
   *       <li>{@code <= 0} -> 1
   *       <li>{@code >= 2^30} -> 2^30
   *     </ul>
   */
  public static int safeFindNextPositivePowerOfTwo(final int value) {
    return value <= 0 ? 1 : value >= 0x40000000 ? 0x40000000 : findNextPositivePowerOfTwo(value);
  }

  private static int free(MemoryRegionCache[] caches, boolean finalizer) {
    if (caches == null) {
      return 0;
    }

    int numFreed = 0;
    for (MemoryRegionCache c : caches) {
      numFreed += free(c, finalizer);
    }
    return numFreed;
  }

  private static int free(MemoryRegionCache cache, boolean finalizer) {
    if (cache == null) {
      return 0;
    }
    return cache.free(finalizer);
  }

  boolean add(
      BufferPool area,
      BufferChunk chunk,
      Buffer nioBuffer,
      long handle,
      int normCapacity,
      BufferPoolArena.bufferArea bufferArea) {
    int sizeIdx = ((BufferPoolArena) area).size2SizeIdx(normCapacity);
    MemoryRegionCache cache = cache(area, sizeIdx, bufferArea);
    if (cache == null) {
      return false;
    }
    return cache.add(chunk, nioBuffer, handle);
  }

  private MemoryRegionCache cache(
      BufferPool area, int sizeIdx, BufferPoolArena.bufferArea bufferArea) {
    switch (bufferArea) {
      case Normal:
        return cacheForNormal(area, sizeIdx);
      case Small:
        return cacheForSmall(area, sizeIdx);
      default:
        throw new Error();
    }
  }

  private MemoryRegionCache cacheForNormal(BufferPool arena, int sizeIdx) {
    // sizeClass == Normal => normCapacity >= pageSize => the shifted value > 0
    int idx = sizeIdx - ((BufferPoolArena) arena).numSmallSubpagePools;
    return cache(normalDirectCaches, idx);
  }

  private MemoryRegionCache cacheForSmall(BufferPool arena, int sizeIdx) {
    return cache(smallSubPageDirectCaches, sizeIdx);
  }

  void trim() {
    trim(normalDirectCaches);
  }

  public boolean allocateNormal(
      BufferPoolArena arena, PooledBuffer buf, int reqCapacity, int sizeIdx) {
    return allocate(cacheForNormal(arena, sizeIdx), buf, reqCapacity);
  }

  private boolean allocate(MemoryRegionCache cache, PooledBuffer buf, int reqCapacity) {
    if (cache == null) {
      // no cache found so just return false here
      return false;
    }
    boolean allocated = cache.allocate(buf, reqCapacity, this);
    if (++allocations >= freeSweepAllocationThreshold) {
      allocations = 0;
      trim();
    }
    return allocated;
  }

  @Override
  protected void finalize() throws Throwable {
    try {
      super.finalize();
    } finally {
      free(true);
    }
  }

  /**
   * Should be called if the Thread that uses this cache is about to exist to release resources out
   * of the cache
   */
  void free(boolean finalizer) {
    // As free() may be called either by the finalizer or by FastThreadLocal.onRemoval(...) we need
    // to ensure
    // we only call this one time.
    if (freed.compareAndSet(false, true)) {
      int numFreed =
          free(smallSubPageDirectCaches, finalizer) + free(normalDirectCaches, finalizer);

      if (numFreed > 0) {
        LOG.d(
            "Freed {} thread-local buffer(s) from thread: {}",
            numFreed,
            Thread.currentThread().getName());
      }

      if (bufferPool != null) {
        ((BufferPoolArena) bufferPool).numThreadCaches.getAndDecrement();
      }
    }
  }

  public boolean allocateSmall(
      BufferPoolArena arena, PooledBuffer buf, int reqCapacity, int sizeIdx) {
    return allocate(cacheForSmall(arena, sizeIdx), buf, reqCapacity);
  }

  /** Cache used for buffers which are backed by TINY or SMALL size. */
  private static final class SubPageMemoryRegionCache extends MemoryRegionCache {
    SubPageMemoryRegionCache(int size) {
      super(size, Small);
    }

    @Override
    protected void initBuf(
        BufferChunk chunk,
        Buffer nioBuffer,
        long handle,
        int reqCapacity,
        PoolThreadCache threadCache,
        PooledBuffer pooledBuffer) {
      chunk.initBufWithSubpage(pooledBuffer, nioBuffer, handle, reqCapacity, threadCache);
    }
  }

  /** Cache used for buffers which are backed by NORMAL size. */
  private static final class NormalMemoryRegionCache extends MemoryRegionCache {
    NormalMemoryRegionCache(int size) {
      super(size, Normal);
    }

    @Override
    protected void initBuf(
        BufferChunk chunk,
        Buffer nioBuffer,
        long handle,
        int reqCapacity,
        PoolThreadCache threadCache,
        PooledBuffer pooledBuffer) {
      chunk.initBuf(pooledBuffer, reqCapacity, threadCache, handle);
    }
  }

  private abstract static class MemoryRegionCache {
    @SuppressWarnings("rawtypes")
    private static final ObjectPool<Entry> RECYCLER =
        ObjectPool.newPool(handle -> new Entry(handle));

    private final int size;
    private final Queue<Entry> queue;
    private int allocations;
    private BufferPoolArena.bufferArea bufferArea;

    MemoryRegionCache(int size, BufferPoolArena.bufferArea bufferArea) {
      this.bufferArea = bufferArea;
      this.size = safeFindNextPositivePowerOfTwo(size);
      queue = JctoolsHelper.mpscArrayQueue2(this.size);
    }

    @SuppressWarnings("rawtypes")
    private static Entry newEntry(BufferChunk chunk, Buffer nioBuffer, long handle) {
      Entry entry = RECYCLER.get();
      entry.chunk = chunk;
      entry.nioBuffer = nioBuffer;
      entry.handle = handle;
      return entry;
    }

    /**
     * Init the {@link Buffer} using the provided chunk and handle with the capacity restrictions.
     */
    protected abstract void initBuf(
        BufferChunk chunk,
        Buffer nioBuffer,
        long handle,
        int reqCapacity,
        PoolThreadCache threadCache,
        PooledBuffer pooledBuffer);

    /** Add to cache if not already full. */
    @SuppressWarnings("unchecked")
    public final boolean add(BufferChunk chunk, Buffer nioBuffer, long handle) {
      Entry entry = newEntry(chunk, nioBuffer, handle);
      boolean queued = queue.offer(entry);
      if (!queued) {
        // If it was not possible to cache the chunk, immediately recycle the entry
        entry.recycle();
      }

      return queued;
    }

    /** Allocate something out of the cache if possible and remove the entry from the cache. */
    public final boolean allocate(PooledBuffer buf, int reqCapacity, PoolThreadCache threadCache) {
      Entry entry = queue.poll();
      if (entry == null) {
        return false;
      }
      initBuf(entry.chunk, entry.nioBuffer, entry.handle, reqCapacity, threadCache, buf);
      entry.recycle();

      // allocations is not thread-safe which is fine as this is only called from the same thread
      // all time.
      ++allocations;
      return true;
    }

    /**
     * Clear out this cache and free up all previous cached {@link BufferChunk}s and {@code
     * handle}s.
     */
    public final int free(boolean finalizer) {
      return free(Integer.MAX_VALUE, finalizer);
    }

    private int free(int max, boolean finalizer) {
      int numFreed = 0;
      for (; numFreed < max; numFreed++) {
        Entry entry = queue.poll();
        if (entry != null) {
          freeEntry(entry, finalizer);
        } else {
          // all cleared
          return numFreed;
        }
      }
      return numFreed;
    }

    /** Free up cached {@link BufferChunk}s if not allocated frequently enough. */
    public final void trim() {
      int free = size - allocations;
      allocations = 0;

      // We not even allocated all the number that are
      if (free > 0) {
        free(free, false);
      }
    }

    @SuppressWarnings({"unchecked", "rawtypes"})
    private void freeEntry(Entry entry, boolean finalizer) {
      BufferChunk chunk = entry.chunk;
      Buffer nioBuffer = entry.nioBuffer;

      if (!finalizer) {
        // recycle now so PoolChunk can be GC'ed. This will only be done if this is not freed
        // because of
        // a finalizer.
        entry.recycle();
      }
      BufferPoolArena arena = (BufferPoolArena) chunk.bufferPool;
      PooledBuffer localBuffer = (PooledBuffer) nioBuffer;
      arena.freeChunk(localBuffer.handle, localBuffer.maxLength, nioBuffer, chunk);
    }

    static final class Entry {
      final ObjectPool.Handle<Entry> recyclerHandle;
      BufferChunk chunk;
      Buffer nioBuffer;
      long handle = -1;

      Entry(ObjectPool.Handle<Entry> recyclerHandle) {
        this.recyclerHandle = recyclerHandle;
      }

      void recycle() {
        chunk = null;
        nioBuffer = null;
        handle = -1;
        recyclerHandle.recycle(this);
      }
    }
  }
}
